# Autor chenfeng
# !/usr/bin/env Python
# coding=utf-8

from PIL import Image
import time

# 创建一个Task
class Taskflow():
    def __init__(self, task, tokenizer, model=None, mode=None, device_id=0, num_return_sequences=8, returnTensor=False, saveimg=True, **kwargs):
        self.task = task
        self.model = model
        self.tokenizer = tokenizer
        self.num_return_sequences = num_return_sequences # 返回图片的数量
        self.returnTensor = returnTensor # 是否返回图片Tensor
        self.saveimg = saveimg # 是否保存图片

    def task_genration(self,prompt):
        self.model.eval()
        # 可选择的超参数
        top_k = 128
        condition_scale = 10.0
        num_return_sequences = self.num_return_sequences

        tokenized_inputs = self.tokenizer(
            prompt,
            return_tensors="pd",
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            max_length=64,
        )

        images = self.model.generate(**tokenized_inputs,
                                     top_k=top_k,
                                     condition_scale=condition_scale,
                                     num_return_sequences=num_return_sequences)
        # 黑白图
        # images = (images.cpu().numpy().clip(0, 1) * 255).astype("uint8")
        # 彩图
        images = (images.cpu().numpy()).astype("uint8")
        images = images.transpose([0, 2, 1, 3, 4]).reshape(-1, images.shape[-3],num_return_sequences * images.shape[-2], images.shape[-1])
        return images[0]

    def __call__(self, *inputs):
        """
        The main work function in the taskflow.
        """
        results = self.task_genration(inputs)
        image = Image.fromarray(results)
        image.show()
        if self.saveimg:
            image.save("generateimages/figure_{}.png".format(int(time.time())))
        if self.returnTensor:
            return results
